Hierarchical Muon: Tiled Newton-Schulz Updates for Efficient Muon Optimization
分层Muon：高效训练的分块Newton-Schulz更新
🧭 高效训练优化器 · 分块Newton-Schulz
通过将动量梯度矩阵分块独立应用Newton-Schulz，大幅降低Muon优化器计算量并实现高效GPU并行
📌 工作参考价值（4 点）
- 在训练含有大型密集矩阵参数的推荐或广告模型时，可将 Muon 优化器的 Newton-Schulz 步改为分块方案，直接降低从 O(r²s) 到 O(HWT) 的单步计算量，同时保持接近全矩阵更新的训练效果。
- 利用独立小方块计算的特性，可以设计 tile-size 相关的 GPU kernel，配合跨层批处理（cross-layer batching），将不同层的同尺寸块合并批量矩阵乘法，显著提升 GPU 吞吐。
- 采用内存受限分块（memory-bounded chunking）和运行时动态调整块尺寸的策略，能灵活适配不同硬件显存限制，避免 OOM 且不牺牲收敛。
- 该方法将矩阵谱交互限定在块内，可作为工程上折衷效果与效率的通用抽象，类似思路可迁移至其他需要 per‑parameter 矩阵变换的优化场景（如 Shampoo 类优化器的分块近似）。